# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from experimental.prophet.graph import graph as graph_mod
from experimental.prophet.graph.shape import Shape
from experimental.prophet.ops.placeholder import placeholder
from experimental.prophet.ops.util import to_duration_int, to_duration_str


def timer(period):
  period = to_duration_int(period)
  assert period >= 10**7  # 10ms
  node_name = '_timer/PeriodicTimer.%s' % to_duration_str(period)

  try:
    node = graph_mod.get_default_graph().get_node(node_name)
    return node.outputs[0]
  except KeyError:
    return placeholder(int,
                       shape=Shape.Scalar(),
                       node_name=node_name,
                       namespace='',
                       attrs={'period': period})


def custom_timer(timestamps):
  assert timestamps
  prev_ts = None
  for ts in timestamps:
    assert prev_ts is None or prev_ts < ts
    assert isinstance(ts, int)
    assert 0 <= ts and ts < 2**63
    prev_ts = ts
  return placeholder(int,
                     shape=Shape.Scalar(),
                     node_name='_timer/CustomTimer',
                     namespace='',
                     attrs={'timestamps': timestamps})


def list_periodic_timers(graph=None):
  graph = graph or graph_mod.get_default_graph()
  timers = []
  for node in graph.nodes:
    if node.name.startswith('_timer/PeriodicTimer.'):
      timers.append((node.outputs[0].instance_name, node.attrs['period']))
  return timers


def list_custom_timers(graph=None):
  graph = graph or graph_mod.get_default_graph()
  timers = []
  for node in graph.nodes:
    if node.name.startswith('_timer/CustomTimer'):
      timers.append((node.outputs[0].instance_name, node.attrs['timestamps']))
  return timers
